import copy
from .resnet import ResNet
from .resnetplusbigmodel import ResNet_BigModel
from .swin import SwinTransformer
from ...utils import BaseModuleBuilder
from .mit import MixVisionTransformer
from .sctnet import SCTNet
from .rmt import VisRetNet
from .repvit import RepViT
from .ddrnet import DDRNet
'''BackboneBuilder'''
class BackboneBuilder(BaseModuleBuilder):
    REGISTERED_MODULES = {
        'ResNet': ResNet, 'SwinTransformer': SwinTransformer, 'SCTNet':SCTNet, 'MixVisionTransformer':MixVisionTransformer,
        'VisRetNet': VisRetNet, 'RepViT':RepViT, 'DDRNet':DDRNet, 'ResNet_BigModel':ResNet_BigModel
    }
    '''build'''
    def build(self, backbone_cfg):
        backbone_cfg = copy.deepcopy(backbone_cfg)
        if 'selected_indices' in backbone_cfg: backbone_cfg.pop('selected_indices')
        return super().build(backbone_cfg)


'''BuildBackbone'''
BuildBackbone = BackboneBuilder().build
